import sys
import cv2
import time
from PyQt5.QtWidgets import QApplication
from ultralytics import YOLO
from multiprocessing import Queue
from pathlib import Path
import json
from threading import Thread
import os


def load_config():
    """加载配置文件"""
    if getattr(sys, 'frozen', False):
        base_dir = Path(sys._MEIPASS)
    else:
        base_dir = Path(__file__).parent

    config_path = base_dir / "config.json"

    try:
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
    except FileNotFoundError:
        sys.exit(f"错误：配置文件 {config_path} 不存在")
    except json.JSONDecodeError:
        sys.exit(f"错误：配置文件格式不正确")

    model_path = Path(config["model_path"])
    if not model_path.exists():
        sys.exit(f"错误：模型文件 {model_path} 不存在")

    video_source = config.get("video_source", "input_video.mp4")
    if not Path(video_source).exists():
        sys.exit(f"错误：视频文件 {video_source} 不存在")

    return {
        "video_source": video_source,
        "model_path": str(model_path),
        "save_dirs": config["save_dirs"],
    }


# 保存未框选的原始帧（队列+线程）
class RawFrameSaver:
    def __init__(self, save_dir="raw_frames"):
        config = load_config()
        self.save_dir = Path(config["save_dirs"]["raw_frames"])
        self.save_dir.mkdir(exist_ok=True)
        self.frame_queue = Queue(maxsize=30)
        self.thread = Thread(target=self._save_worker, daemon=True)
        self.thread.start()

    def add_frame(self, frame, timestamp_str):
        if self.frame_queue.qsize() < 30:
            self.frame_queue.put((frame.copy(), timestamp_str))

    def _save_worker(self):
        while True:
            if self.frame_queue.empty():
                time.sleep(0.1)
                continue
            frame, timestamp_str = self.frame_queue.get()
            filename = f"{timestamp_str}.jpg"
            cv2.imwrite(str(self.save_dir / filename), frame)


# 保存带框的完整帧（直接保存，无需队列）
class FrameSaver:
    def __init__(self, save_dir="full_frames"):
        config = load_config()
        self.save_dir = Path(config["save_dirs"]["full_frames"])
        self.save_dir.mkdir(exist_ok=True)

    def save_full_frame(self, frame, timestamp_str):
        filename = f"{timestamp_str}.jpg"
        save_path = self.save_dir / filename
        cv2.imwrite(str(save_path), frame)


# 图像保存模块（保存裁剪的目标）
class ImageSaver:
    def __init__(self, save_dir="detected_objects"):
        config = load_config()
        self.save_dir = Path(config["save_dirs"]["detected_objects"])
        self.save_dir.mkdir(exist_ok=True)
        self.frame_counter = 0
        self.last_save_time = time.time()
        self.min_interval = 0.5  # 裁剪图防抖间隔

    def save_cropped_image(self, frame, box, class_name, conf):
        current_time = time.time()
        if current_time - self.last_save_time < self.min_interval:
            return

        x1, y1, x2, y2 = map(int, box)
        cropped = frame[y1:y2, x1:x2]
        if cropped.size == 0:
            return

        timestamp = time.strftime("%Y%m%d_%H%M%S")
        filename = f"{timestamp}_{class_name}_{conf:.2f}_{self.frame_counter:04d}.jpg"
        save_path = self.save_dir / filename

        try:
            ret, buffer = cv2.imencode('.jpg', cropped, [cv2.IMWRITE_JPEG_QUALITY, 90])
            if ret:
                buffer.tofile(str(save_path))
                self.frame_counter += 1
                self.last_save_time = current_time
        except Exception as e:
            print(f"保存图像时发生错误：{str(e)}")


# 推理处理类（优化版：所有检测图都保存）
class InferenceProcessor:
    def __init__(self, model_path):
        self.model = YOLO(model_path)
        self.saver = ImageSaver()
        self.frame_saver = FrameSaver()
        self.raw_saver = RawFrameSaver()
        self.class_colors = {
            "hel": (0, 255, 0),      # 绿色
            "nohel": (0, 0, 255),    # 红色
            "wheel": (255, 0, 0),    # 蓝色
        }

    def process_frame(self, frame, timestamp_str):
        height, width, _ = frame.shape
        base_height = 720.0
        font_scale = max(0.6, height / base_height * 0.8)
        font_thickness = max(1, int(font_scale * 1.5))

        raw_frame = frame.copy()
        detected_any = False

        # ✅ 全局推理：只做一次
        results = self.model(frame, imgsz=640, conf=0.5, iou=0.45)

        for result in results:
            boxes = result.boxes.xyxy.cpu().numpy()
            confidences = result.boxes.conf.cpu().numpy()
            class_ids = result.boxes.cls.cpu().numpy().astype(int)

            for i, (xyxy, conf, cls_id) in enumerate(zip(boxes, confidences, class_ids)):
                if conf < 0.5:
                    continue
                class_name = result.names[cls_id]
                x1, y1, x2, y2 = map(int, xyxy)

                # --- 绘制边界框和标签 ---
                color = self.class_colors.get(class_name, (255, 255, 255))

                # 边界框
                cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness=3)

                # 标签文本
                label = f"{class_name} {conf:.2f}"
                font = cv2.FONT_HERSHEY_DUPLEX
                font_scale = 0.7
                font_thickness = 2
                (text_width, text_height), _ = cv2.getTextSize(label, font, font_scale, font_thickness)

                text_x, text_y = x1, max(0, y1 - 10)
                cv2.rectangle(frame,
                              (text_x, text_y - text_height - 8),
                              (text_x + text_width + 6, text_y + 4),
                              color, -1)
                cv2.putText(frame, label,
                            (text_x + 3, text_y - 4),
                            font,
                            font_scale,
                            (255, 255, 255),
                            font_thickness,
                            cv2.LINE_AA)

                # 保存裁剪图像
                self.saver.save_cropped_image(frame.copy(), [x1, y1, x2, y2], class_name, conf)
                detected_any = True

        # ✅ 检测到目标则保存原始图和带框图
        if detected_any:
            self.raw_saver.add_frame(raw_frame, "detect_" + timestamp_str)
            self.frame_saver.save_full_frame(frame.copy(), "detect_" + timestamp_str)

        return frame


# 主程序入口
if __name__ == "__main__":
    app = QApplication(sys.argv)

    config = load_config()

    # 创建保存目录
    for dir_path in config["save_dirs"].values():
        Path(dir_path).mkdir(exist_ok=True)

    # 初始化推理处理器
    processor = InferenceProcessor(
        model_path=config["model_path"]
    )

    # 打开视频文件
    video_path = config["video_source"]
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        sys.exit(f"无法打开视频文件：{video_path}")

    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    print(f"视频 FPS: {fps}, 总帧数: {total_frames}")

    frame_idx = 0
    while True:
        frame_idx += 1
        if frame_idx % 2 == 0: continue
        ret, frame = cap.read()
        if not ret:
            print("视频读取完成或出错。")
            break

        # 生成时间戳（精确到毫秒）
        current_time_ms = int((frame_idx / fps) * 1000)
        timestamp_str = f"{Path(video_path).stem}_{current_time_ms:06d}"

        print(f"正在处理帧 {frame_idx + 1}/{total_frames} | 时间: {current_time_ms}ms")
        processed_frame = processor.process_frame(frame, timestamp_str)

        # 可选：显示实时画面（调试用）
        # cv2.imshow("Detection", processed_frame)
        # if cv2.waitKey(1) & 0xFF == ord('q'):
        #     break



    cap.release()
    cv2.destroyAllWindows()
    print("视频处理完成，程序退出。")
    sys.exit(0)